import os

from langchain import PromptTemplate
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import TextLoader
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from zhipuai import ZhipuAI

# 从环境变量中获取 API 密钥
api_key = os.getenv("OPENAI_API_KEY")


def ask_questions(rag: str) -> str:
    # 初始化 LLM
    llm = ChatOpenAI(
        temperature=0.95,
        model="glm-4",
        openai_api_key=api_key,
        openai_api_base="https://open.bigmodel.cn/api/paas/v4/",
        streaming=True
    )
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    # 获取当前文件所在的目录 (相对于当前脚本文件)
    current_dir = os.path.dirname(__file__)
    # 构建相对路径
    parent_dir = os.path.dirname(current_dir)
    file_path = os.path.join(parent_dir, 'resource', 'document.md')
    # 加载文档并分割成块
    loader = TextLoader(file_path)

    documents = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    docs = text_splitter.split_documents(documents)

    # 创建向量存储
    vectorstore = FAISS.from_documents(docs, embeddings)

    # 创建检索器
    retriever = vectorstore.as_retriever()

    # 自定义 Prompt 模板
    prompt_template = """
    根据以下内容回答问题:
    {context}

    问题: {question}
    答案:"""
    prompt = PromptTemplate(input_variables=["context", "question"], template=prompt_template)

    # 创建 RAG 链
    rag_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever,
                                            chain_type_kwargs={"prompt": prompt})

    # 使用 RAG 链进行查询
    return rag_chain.run(rag)


def text_to_vector(text: str) -> str:
    client = ZhipuAI(api_key=api_key)  # 填写您自己的APIKey
    response = client.chat.completions.create(
        model="glm-4-0520",  # 填写需要调用的模型编码
        messages=[
            {"role": "user", "content": "作为一名营销专家，请为智谱开放平台创作一个吸引人的slogan"},
            {"role": "assistant", "content": "当然，为了创作一个吸引人的slogan，请告诉我一些关于您产品的信息"},
            {"role": "user", "content": "智谱AI开放平台"},
            {"role": "assistant", "content": "智启未来，谱绘无限一智谱AI，让创新触手可及!"},
            {"role": "user", "content": "创造一个更精准、吸引人的slogan"}
        ],
        stream=True
    )
    str = ""
    for chunk in response:
        str += chunk.choices[0].delta.content
    print(str)


if __name__ == '__main__':
    # text_to_vector('你好')
    print(ask_questions('审批流如何创建'))
